#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from concurrent import futures
import grpc
import signal
import sys
import json
from typing import Iterable
import grpc_pb2
import grpc_pb2_grpc
from answer import complete_answer
from query import query_db
from analyze import analyze

class AnswerServicer(grpc_pb2_grpc.Go_Py_AnswerServicer):
    def GetAnswer(self, request, context) -> Iterable[grpc_pb2.AnswerResponse]:
        """流式问答服务实现（适配您的生成器）"""
        try:
            print(f"[Answer] 收到提示词: {request.cue}")
            
            # 获取生成器实例
            answer_generator = complete_answer(request.cue)
            
            # 第一轮快速响应（避免冷启动延迟）
            first_chunk = next(answer_generator, None)
            if first_chunk:
                yield grpc_pb2.AnswerResponse(
                    text=first_chunk,
                    done=False
                )
            
            # 流式传输剩余内容
            for chunk in answer_generator:
                #每次响应前检查连接状态（核心优化点）
                if not context.is_active():
                    print("客户端断开，终止生成")  # 日志记录
                    answer_generator.close()  # 主动释放资源
                    break
                yield grpc_pb2.AnswerResponse(
                    text=chunk,
                    done=False
                )
            
            # 最终结束标记
            yield grpc_pb2.AnswerResponse(
                text="",
                done=True
            )
            
        except StopIteration:
            # 生成器正常结束
            yield grpc_pb2.AnswerResponse(
                text="",
                done=True
            )
        except Exception as e:
            # 标准化错误状态码（Go兼容关键）
            context.set_code(grpc.StatusCode.INTERNAL)  
            context.set_details(f"服务端错误: {str(e)}")  # 传递详情
            print(f"[Error] 流式处理异常: {str(e)}")
            yield grpc_pb2.AnswerResponse(
                text=f"服务错误: {str(e)}",
                done=True
            )

class GraphDataServicer(grpc_pb2_grpc.Go_Py_GraphDataServicer):
    def GetGraphData(self, request, context):
        """图数据查询服务实现"""
        try:
            print(f"[Graph] Question: {request.question}")
            
            # 替换为您的图数据库查询逻辑
            answer,cypher  = query_db(request.question)
            print("type=",type(answer))
            print("answer=",answer)
            return grpc_pb2.GraphDataResponse(answer=answer,cypher_statement=cypher)
            
        except Exception as e:
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(str(e))
            return grpc_pb2.GraphDataResponse(answer=str(e))
        
class AnalyzeServicer(grpc_pb2_grpc.Go_Py_AnalyzeServicer):
    def Analyze(self, request, context):
        """解析文件生成图数据库数据服务实现"""
        try:
            success,label = analyze(request.content)
            print("type=",type(success))
            print("answer=",success)
            return grpc_pb2.AnalyzeResponse(success=success,label=label)
            
        except Exception as e:
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(str(e))
            return grpc_pb2.AnalyzeResponse(answer=str(e))

def serve():
    # 配置高性能服务器
    server = grpc.server(
        futures.ThreadPoolExecutor(max_workers=100),
        options=[
            ('grpc.max_concurrent_streams', 1000),
            ('grpc.max_receive_message_length', 100 * 1024 * 1024)
        ]
    )
    
    # 注册所有服务
    grpc_pb2_grpc.add_Go_Py_AnswerServicer_to_server(AnswerServicer(), server)
    grpc_pb2_grpc.add_Go_Py_GraphDataServicer_to_server(GraphDataServicer(), server)
    grpc_pb2_grpc.add_Go_Py_AnalyzeServicer_to_server(AnalyzeServicer(),server)
    
    # 启动服务
    server.add_insecure_port('[::]:50052')
    server.start()
    print("gRPC服务已启动 (端口:50052)\n"
          "支持服务:\n"
          "- Go_Py_Answer/GetAnswer(stream)\n"
          "- Go_Py_GraphData/GetGraphData")
    
    # 优雅退出处理
    def shutdown(sig, frame):
        print(f"\n终止信号: {sig}, 关闭服务...")
        server.stop(grace=5)
        sys.exit(0)
    
    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)
    server.wait_for_termination()

if __name__ == '__main__':
    serve()